1
從線程轉向程式實例
AI023Lesson 4
00:00

在 Triton 中,執行的基本單位從 CUDA 標量線程轉移至 程式實例。這代表了 GPU 線程塊的一種抽象,其中單一實例可同時處理一個向量化「區塊」的元素。

1. 程式實例身分

每個執行單元皆透過以下方式取得其身分 pid = tl.program_id(axis=0)。想像一輛 倉儲叉車 (程式實例)搬運一個 托盤 (區塊)共 128 個箱子,與單一工人(CUDA 線程)僅搬一箱形成對比。

2. Triton 與 PyTorch 張量的對比

理解語義差異對於記憶體管理至關重要:

  • PyTorch 張量: 位於主機端的 Python 物件,包裝了顯存儲存空間、步長與元資料。
  • Triton 張量: 編譯器層級的物件,代表位於 暫存器或 SRAM的值或指標。
PyTorch 視圖
指向連續全域記憶體的 Python 物件。
Triton 視圖
編譯器暫存器內的二維/一維資料區塊。

3. SPMD 特性

Triton 遵循 單一程式、多重資料(SPMD) 流程。每個程式實例都執行 完全相同的 程式碼。分歧僅當邏輯利用 pid 來計算特定的記憶體偏移量時才會發生。

main.py
TERMINALbash — 80x24
> Ready. Click "Run" to execute.
>